import os
import numpy as np
from sklearn.decomposition import PCA
import cv2
import torch

def embedding_pca(embedding, n_components=3,as_rgb=True):
    pca = PCA(n_components=n_components)
    embed_dim = embedding.shape[0]
    shape = embedding.shape[1:]
    
    embed_flat= embedding.reshape(embed_dim, -1).T
    embed_flat = pca.fit_transform(embed_flat).T
    embed_flat =embed_flat.reshape((n_components,)+shape)
    
    
    if as_rgb:
        embed_flat = 255*(embed_flat-embed_flat.min())/np.ptp(embed_flat)
        embed_flat = np.transpose(embed_flat, (1,2,0))
        embed_flat = embed_flat.astype(np.uint8)
    return embed_flat

def main():
    embedding = torch.rand(100, 100, 100)
    embed_flat =embedding_pca(embedding, n_components=3, as_rgb=True)
    print(embed_flat.shape)
    cv2.imwrite("embed_flat.png", embed_flat)


if __name__ == "__main__":
    main()
    